#!/usr/bin/python
# -*- coding: utf-8 -*-

import serial
import time

from elftools.elf.elffile import ELFFile
from elftools.elf.constants import SH_FLAGS

import sparc_dis
from bm3823_def import *
from bm3823_ddr import *

class DSU_Comm:
	def __init__(self):
		self.init_baudrate = 38400
		self.curr_baudrate = self.init_baudrate

		self.serialPort_init = False
		self.dsu_comm_init = False
		self.prog_running  = False

#		self.getmode_error = False
#		self.model = ""
#		self.scaler = ""
#		self.in_stream = None
#		self.out_stream = None
#		self.serialPort_init = False
#		self.serial_port = None
#		self.serialPort_in_thread = None
#		self.serialPort_out_thread = None

#		self.port_name = port_name

	def connect(self, port_name, baudrate=38400):
		"""连接串口"""
		try:
			self.port_name = port_name
			self.curr_baudrate = baudrate

			self.serial_port = serial.Serial(
				port=self.port_name,
				baudrate=self.curr_baudrate,
				bytesize=serial.EIGHTBITS,
				parity=serial.PARITY_NONE,
				stopbits=serial.STOPBITS_ONE,
				timeout=1
			)
			self.in_stream = self.serial_port
			self.out_stream = self.serial_port

			self.in_stream.flushInput()
			self.out_stream.flushOutput()

			self.serialPort_init = True
			
			return True
		except Exception as e:
			print(f"串口连接错误: {e}")
			self.serialPort_init = False
			return False

	def close_serial_port(self):
		"""关闭串口"""
		self.serialPort_init = False
		try:
			if self.serial_port:
				self.serial_port.close()
		except Exception as e:
			print(f"关闭串口错误: {e}")

	def dsu_uart_tx(self, dsu_str):
		bytes_count = len(dsu_str)
		tq = 1.0 / self.curr_baudrate
#		sleep_time = tq*20*bytes_count + 0.005
		sleep_time = tq*20*bytes_count

		self.in_stream.flushInput()
		self.out_stream.flushOutput()
		self.out_stream.write(dsu_str)
		time.sleep(sleep_time)

	def dsu_uart_rx(self, bytes_count):
		tq = 1.0 / self.curr_baudrate
		sleep_time  = tq*20*bytes_count
		retry_count = 0

		time.sleep(sleep_time+0.005)

		while(self.in_stream.inWaiting() != bytes_count):
			retry_count += 1
			if(retry_count > 5):
				print("READ UART FAIL!")
				break
			time.sleep(0.005)

		res = self.in_stream.read(bytes_count)

		return res

	def dsu_send_55(self):
		dsu_str = b'\x55' * 32

		self.dsu_uart_tx(dsu_str)
#		return dsu_str

	def dsu_reg_read(self, reg_start, dword_count):
		if(not isinstance(reg_start, int)) or (not isinstance(dword_count, int)):
			raise TypeError("TYPE ERROR!")
		if((reg_start % 4) != 0):
			raise ValueError("ADDR NOT ALIGNED!")

		if(dword_count > 0x40):
			raise ValueError("Too much words to read!")

		dsu_str_len = 1 + 4

		dsu_cmd = 0x80 + dword_count - 1
		dsu_str = dsu_cmd.to_bytes(1, byteorder="big", signed=False) + reg_start.to_bytes(4, byteorder="big", signed=False)

#		print(dsu_str)
		self.dsu_uart_tx(dsu_str)
		res = self.dsu_uart_rx(dword_count*4)
#		print(res)

		reg_list = []
		for _ in range(int(len(res)/4)):
			reg_tmp = res[_*4:(_+1)*4]
			reg_val = int.from_bytes(reg_tmp, byteorder='big', signed=False)
			reg_list.append(reg_val)

		return reg_list

	def dsu_reg_write(self, reg_start, dword_list):
		if(not isinstance(reg_start, int)) or (not isinstance(dword_list, list)):
			raise TypeError("TYPE ERROR!")
		if((reg_start % 4) != 0):
			raise ValueError("ADDR NOT ALIGNED!")

		dword_count = len(dword_list)
		if(dword_count > 0x40):
			raise ValueError("Too much words to write!")

		dsu_str_len = 1 + 4 + (dword_count * 4)

		dsu_cmd = 0xC0 + dword_count - 1
		dsu_str = dsu_cmd.to_bytes(1, byteorder="big", signed=False) + reg_start.to_bytes(4, byteorder="big", signed=False)

		for _ in dword_list:
			dsu_str = dsu_str + _.to_bytes(4, byteorder="big", signed=False)

		if(dsu_str_len != len(dsu_str)):
			raise ValueError("Check dword_list!")

		self.dsu_uart_tx(dsu_str)
#		return dsu_str

	def dsu_mem_write(self, mem_start, data_bytes):
		if(not isinstance(mem_start, int)) or (not isinstance(data_bytes, bytes)):
			raise TypeError("TYPE ERROR!")
		if((mem_start % 4) != 0):
			raise ValueError("ADDR NOT ALIGNED!")

		bytes_count = len(data_bytes)
		if((bytes_count % 4) != 0):
			raise ValueError("MEM NOT ALIGNED!")

		dword_count = bytes_count >> 2
		if(dword_count > 0x40):
			raise ValueError("Too much words to write!")

		dsu_str_len = 1 + 4 + bytes_count

		dsu_cmd = 0xC0 + dword_count - 1
		dsu_str = dsu_cmd.to_bytes(1, byteorder="big", signed=False) + mem_start.to_bytes(4, byteorder="big", signed=False)
		dsu_str = dsu_str + data_bytes

		if(dsu_str_len != len(dsu_str)):
			raise ValueError("Check data_bytes!")

		self.dsu_uart_tx(dsu_str)
#		return dsu_str

	def bm3823_flush_idcache(self):
		self.dsu_reg_write(0x80000014, [0x00E00000])

	def bm3823_internal_sram_init(self):
		self.dsu_reg_write(0x80000030, [0xC0000000])

	def bm3823_dsu_init_seq(self):
		self.dsu_reg_write(0x90000000, [0x000075FA, 0x031FF1FF, 0x00000000, 0x00000000])
		self.dsu_reg_write(0x90080040, [0x00000000])
		self.dsu_reg_write(0x90080060, [0x00000000, 0x00000000, 0x00000000, 0x00000000, 0x00000000, 0x00000000, 0x00000000, 0x00000000])
		self.dsu_reg_write(0x901E0000, [0x00000000])
		self.dsu_reg_write(0x81000004, [0xFD27D06F, 0x00208000])
		self.dsu_reg_write(0x80000014, [0x00E1000F])
		self.dsu_reg_write(0x80000030, [0xC0000000])

#		self.dsu_reg_write(0x90080040, [0x00000000])
#		self.dsu_reg_write(0x90080060, [0x00000000, 0x00000000, 0x00000000, 0x00000000, 0x00000000, 0x00000000, 0x00000000, 0x00000000])
#		self.dsu_reg_write(0x90000000, [0x000075FA])
#		self.dsu_reg_write(0x901E0000, [0x00000000])
#		self.dsu_reg_write(0x80000014, [0x00E1000F])
#		self.dsu_reg_write(0x90000008, [0x00000000])
##		self.dsu_reg_write(0x90000004, [0x031FF1FF])
#		self.dsu_reg_write(0x90000004, [0x03000000])
#		self.dsu_reg_write(0x81000004, [0x7C478E6F])
#		self.dsu_reg_write(0x80000030, [0xC0000000])

	def bm3823_ddr_init_seq(self):
		self.dsu_reg_write(BM3823_DDR_CTL_GRP0_ADDR, BM3823_CTL_GRP0)
		self.dsu_reg_write(BM3823_DDR_CTL_GRP1_ADDR, BM3823_CTL_GRP1)
		self.dsu_reg_write(BM3823_DDR_CTL_GRP2_ADDR, BM3823_CTL_GRP2)

		self.dsu_reg_write(BM3823_DDR_PHY_GRP0_ADDR, BM3823_PHY_GRP0)
		self.dsu_reg_write(BM3823_DDR_PHY_GRP1_ADDR, BM3823_PHY_GRP1)
		self.dsu_reg_write(BM3823_DDR_PHY_GRP2_ADDR, BM3823_PHY_GRP2)

		ctl_000_data  = self.dsu_reg_read(BM3823_DDR_CTL_000_ADDR, 1)[0]
		ctl_000_data |= 0b1
		self.dsu_reg_write(BM3823_DDR_CTL_000_ADDR, [ctl_000_data])

		while(True):
			ctl_066_data  = self.dsu_reg_read(BM3823_DDR_CTL_066_ADDR, 1)[0]
			if((ctl_066_data & 0x200) != 0):
				break
			time.sleep(0.005)

		print("DDR2 INIT DONE!")


	def bm3823_chip_id_get(self):
		model = ""
		reg = self.dsu_reg_read(0x8100010C, 1)
		if(reg == []):
			print("READ FAIL!")
			return ""
		chip_id = reg[0] & 0x0000000F

		if chip_id == 0x3:
			model = "BM3803"
		elif chip_id == 0x6:
			model = "BM3110"
		elif chip_id == 0x8:
			model = "BM3803"
		elif chip_id == 0x9:
			model = "BM3107"
		elif chip_id == 0xA:
			model = "BM3823"

		return model

	def bm3823_dsu_connect(self):
		retry_count = 0
		res = False
		while(True):
			self.dsu_send_55()
			model = self.bm3823_chip_id_get()
			if(model != ""):
				print("CONNECT BOARD SUCCESS!")
				res = True
				break
			retry_count += 1
			time.sleep(0.010)
			if(retry_count > 20):
				print("CONNECT BOARD FAIL!")
				res = False
				break

		self.dsu_comm_init = res
		return res

	def bm3823_freq_get(self):
		reg = self.dsu_reg_read(0x800000CC, 1)
		if(reg == []):
			print("READ FAIL!")
			return ""
		scaler = reg[0] & 0x0003FFFF
		print("bm3823_freq_get:scaler:0x%04X" % scaler)

		frequency = (scaler * 10 + 5) * 8.0
		frequency *= int(self.curr_baudrate)
		frequency /= 1.0e7
		frequency = round(frequency * 100) / 100

		return (frequency, scaler)

	def bm3823_dsu_scaler_calc(self, freq, baudrate):
		scaler = ((freq * 10.0) / (baudrate * 8.0) - 5.0) / 10.0
		scaler = int(round(scaler)) & 0x0003FFFF
#		print("bm3823_dsu_scaler_set:scaler:0x%04X" % scaler)
		return scaler

	def bm3823_dsu_scaler_set(self, freq, baudrate):
		scaler = self.bm3823_dsu_scaler_calc(freq, baudrate)

		self.dsu_reg_write(0x800000CC, [scaler])

		self.serial_port.close()
		self.curr_baudrate = baudrate
		self.serial_port.baudrate = self.curr_baudrate
		self.serial_port.open()
		self.serial_port.flushInput()
		self.serial_port.flushOutput()

	def bm3823_bin_loader(self, mem_start, bin_buffer, speed_flag=False):
		if((len(bin_buffer) % 4) != 0):
			bin_buffer = bin_buffer + (b'\x00' * (4 - (len(bin_buffer) % 4)))
		bytes_count = len(bin_buffer)

		cycles = int(bytes_count / 256)
		if((bytes_count % 256) != 0):
			cycles = cycles + 1

#		print("Load %d Bytes to 0x%08X" % (bytes_count, mem_start))
		t0 = time.monotonic()

		for i in range(cycles):
			j = i * 256
			k = j + 256
			if k > bytes_count:
				k = bytes_count

			curr_addr = mem_start+j

			if(i % 16 == 0):
				print("Loading to 0x%08X, %8d bytes to load" % (curr_addr, (bytes_count - j)))

			self.dsu_mem_write(curr_addr, bin_buffer[j:k])

		if(speed_flag == True):
			t1 = time.monotonic()
			delta_t = t1-t0
			load_speed = (bytes_count / delta_t / 1024.0)
			print("DONE, time:%.3fs, %.3f kBytes/s" % (delta_t, load_speed))

		return bytes_count

	def bm3823_elf_loader(self, elf_file_dir, speed_flag=False):
		f = open(elf_file_dir, "rb")
		elffile = ELFFile(f)

		if(not elffile):
			print("File is not ELF!")
			f.close()
			return

		header = elffile.header
		entry  = header.e_entry
		sects  = header.e_shnum

		total_size = 0

		t0 = time.monotonic()
		for section in elffile.iter_sections():
			sh_flags  = section['sh_flags']
			sh_type   = section['sh_type']
			sh_addr   = section['sh_addr']
			sect_size = section.data_size

#			print("SH_FLAGS: 0x%04X, SH_TYPE:%s, ADDR:0x%08X" % (sh_flags, sh_type, sh_addr))

			if ((sh_flags & (SH_FLAGS.SHF_EXECINSTR | SH_FLAGS.SHF_ALLOC | SH_FLAGS.SHF_WRITE)) and (sh_type != 'SHT_NOBITS')):
				sect_data = section.data()
				self.bm3823_bin_loader(sh_addr, sect_data, speed_flag=False)
				total_size += sect_size

		if(speed_flag == True):
			t1 = time.monotonic()
			delta_t = t1-t0
			load_speed = (total_size / delta_t / 1024.0)
			print("DONE, time:%.3fs, %.3f kBytes/s" % (delta_t, load_speed))

		f.close()

		if(total_size == 0):
			return -2

		print("ELF Load DONE, Entry point addr: 0x%08X size = %8d bytes" % (entry, total_size))

		return entry

	def bm3823_dsu_trace_disable(self):
		DCR = self.dsu_reg_read(0x90000000, 1)[0]

		# Unset TE bit to disable trace
		DCR &= 0xFFFFFFFE

		self.dsu_reg_write(0x90000000, [DCR])

	def bm3823_dsu_set_debug_mode(self):
		DCR = self.dsu_reg_read(0x90000000, 1)[0]

		# Unset TE bit to disable trace
		DCR &= 0xFFFFFFFE
		# Set BN bit to enable debug mode
		DCR |= 0x00000080

		self.dsu_reg_write(0x90000000, [DCR])
		self.prog_running = False

	def bm3823_dsu_set_normal_mode(self):
		DCR = self.dsu_reg_read(0x90000000, 1)[0]

		# Set TE bit to enable trace
		DCR |= 0x00000001
		# Unset BN bit to resume normal execution
		DCR &= 0xFFFFFF7F

		self.dsu_reg_write(0x90000000, [DCR])
		self.prog_running = True

	def bm3823_running_check(self):
		DCR = self.dsu_reg_read(0x90000000, 1)[0]
		dissambler = sparc_dis.SPARC_DISASSEMBLER()
		if((DCR & 0x00000080) != 0):
			self.bm3823_dsu_set_debug_mode()

			reg_group = self.dsu_reg_read(0x90080000, 8)

			PC  = reg_group[4]
			nPC = reg_group[5]
			DTR = reg_group[7]

			DTR_TT = (DTR >>  4) & 0xFF
			DTR_EM = (DTR >> 12) & 0x1

			print("DTR: 0x%08X" % (DTR))
			if(DTR_EM != 0):
				trap_value = self.dsu_reg_read(PC, 1)[0]
#				self.bm3823_dsu_trace_disable()
				print("Break on 0x%08X: %08X, trap:0x%02X, nPC:0x%08X" % (PC, trap_value, DTR_TT, nPC))
				print("%08X: %s" % (PC, dissambler.sparc_opcode_dis(PC, trap_value)))
			else:
				print("trap:0x%02X" % (DTR_TT))

	def bm3823_go_nobreak(self, go_addr, wait_flag=False):
		if(not isinstance(go_addr, int)):
			raise TypeError("TYPE ERROR!")
		if((go_addr % 4) != 0):
			raise ValueError("ADDR NOT ALIGNED!")

		self.bm3823_dsu_set_debug_mode()

		reg_group = self.dsu_reg_read(0x90080000, 8)

		reg_group[4] = go_addr
		reg_group[5] = go_addr + 4

		print("Start free run from 0x%08X" % (go_addr))
		self.dsu_reg_write(0x90080000, reg_group[0:7])

		self.bm3823_dsu_set_normal_mode()

if __name__ == "__main__":
	bin_file = open("../demo/uart_test/bin/uart_test.bin", "rb")
	bin_file_cont = bin_file.read()
	bin_file.close()

	comm = DSU_Comm()

	BM3823_OSC_FREQ      = 25000000
	BM3823_PLL_MIN       = 0b011
	BM3823_PLL_DIN       = 0b010
	BM3823_PLL_DDR_DIN   = 0b110

	BM3823_MAIN_PLL_FREQ = BM3823_OSC_FREQ * BM3823_PLL_MIN * 16
	BM3823_SYSTEM_FREQ   = BM3823_MAIN_PLL_FREQ / (2 * BM3823_PLL_DIN)
	BM3823_DDR_FREQ      = BM3823_MAIN_PLL_FREQ / BM3823_PLL_DDR_DIN
	BM3823_MAC_FREQ      = BM3823_OSC_FREQ * 20
	BM3823_UART_FREQ     = BM3823_SYSTEM_FREQ / 3

	BM3823_SYSTEM_FREQ_M = BM3823_SYSTEM_FREQ / 1e6

	if comm.connect(port_name="COM4", baudrate=38400):
		print("串口连接成功")
		if comm.bm3823_dsu_connect():
			model = comm.bm3823_chip_id_get()
			print("芯片型号:%s" % (model))

			freq, scaler_get  = comm.bm3823_freq_get()
			print("芯片主频设置值:%.3f MHz, 计算值:%.3f MHz" % (BM3823_SYSTEM_FREQ_M, freq))

#			if(abs(BM3823_SYSTEM_FREQ_M - freq) > 2.0):
#				print("芯片主频设置值:%f MHz, 计算值:%f MHz" % (BM3823_SYSTEM_FREQ_M, freq))

			comm.bm3823_dsu_scaler_set(BM3823_SYSTEM_FREQ, 460800)
			time.sleep(0.005)

			model = comm.bm3823_chip_id_get()
			print("芯片型号:%s" % (model))
			freq, scaler_get  = comm.bm3823_freq_get()
			print("芯片主频设置值:%.3f MHz, 计算值:%.3f MHz" % (BM3823_SYSTEM_FREQ_M, freq))

			comm.bm3823_dsu_init_seq()
#			comm.bm3823_uart_init_seq()

#			comm.bm3823_bin_loader(0xA0000000, bin_file_cont, speed_flag=True)
#			comm.bm3823_go_nobreak(0xA0000000)

			entry = comm.bm3823_elf_loader("../demo/uart_test/bin/uart_test.elf", speed_flag=True)
			if(entry > 0):
				comm.bm3823_go_nobreak(entry)

			comm.close_serial_port()

